import altair as alt
import os, psutil, time, watermark

START_TIME = time.perf_counter()
print(f"Loading functions...", end='', flush=True)

# * altair theme

def theme_sci():
    config = {
        "config": {
            "view": {"continuousWidth": 400,
                     "continuousHeight": 300,
                     # border
                     'stroke':False,
                     # "stroke":'gray',
                     # "strokeWidth":.5,
                     },
            "mark": {
                "color": "black",
                'strokeWidth' : 1.7,
                'rule' : {
                    "strokeDash": [4, 2],
                    "strokeWidth": 1.7   
                },
            },
            "circle": {
                # 'size' : 40,
                "strokeOpacity": 1,
                "fillOpacity"  : 1,
                "strokeWidth"  : .01,
            },
            "line": {
                # 'size' : 40,
                "strokeOpacity": 1,
                "fillOpacity"  : 1,
                "strokeWidth"  : 1.2,
            },
            "point": {
                'size' : 25,
                "strokeOpacity": 1,
                "fillOpacity"  : 1,
                "strokeWidth"  : .01,
            },
            'bar':{                
                # fillOpacity:.8,
                "stroke":'white',
                "strokeWidth":.5,
                "yOffset":0
            },
            "scale":{'zero':False,
                     # "pointPadding":1,
                     # "continuousPadding":1
                     # 'bandPaddingInner':.5,
                     # 'bandPaddingOuter':.5,
                     # 'pointPadding':.5
                     },
            "axisX":{},
            "axisY":{},
            "axis": {
                "labelFontSize": 8,
                "titleFontSize": 10,
                "labelColor": "black",
                "titleColor": "black",
                "labelAngle":0,

                # ticks
                "tickColor" : 'black',
                "domainColor": "black",
                "domainWidth": .7,

                # grid
                'grid':False,
                'gridColor':'lightgray',
                'gridOpacity':.7,
                'gridDash':[2,2],
                "gridWidth":.5,
            },
            'range' :{
                "category": {'scheme':'greys'},
                'continuous' : {'scheme':'greys'},
                "symbol": ['circle', 'square', 'triangle', 'triangle-down', 'diamond', 'cross'],
            },
            'legend': {
                'orient' : 'top',
                'titleFontSize':10,

                'labelFontSize':9,
                # 'symbolWidth':200,
                # "symbolHeight": 50,
                # "symbolOffsetdict":100,
            },
            # facet alignment
            'facet':{
                # space between facets
                'spacing': 2
            },
            "header": {
                "labelFontSize": 9,  # Optional: Adjust font size
                "labelFontWeight": "bold",
                "labelAnchor": "start",
                "labelOrient": "top",
                "labelAlign": "left",
                "labelPadding": 0,

                "titleFontSize": 10,  # Optional: Adjust title size
                "titleAlign": "left",
                "titleAnchor": "start",
                "title"      :'',
            },
            "headerRow": {
                "labelFontSize": 9,  # Optional: Adjust font size
                "labelFontWeight": "bold",
                "labelOrient": "right",  # Place row headers on the right
                "labelAlign": "center",
                "labelAnchor": "middle",
                "labelPadding": 0,

                "titleFontSize": 10,  # Optional: Adjust title size
                "titleOrient": "right",  # Place row titles on the right
                "titleAlign": "center",
                "titleAnchor": "middle",
            },
            "headerColumn": {
                "labelFontSize": 9,  # Optional: Adjust font size
                "labelFontWeight": "bold",
                "labelAnchor": "start",
                "labelOrient": "top",
                "labelAlign": "left",
                "labelPadding": 0,

                "titleFontSize": 10,  # Optional: Adjust title size
                "titleAlign": "left",
                "titleAnchor": "middle",
            },
            # plot title
            "title": {
                "fontSize": 14,          # Set the title font size
                # "font": "Arial",         # Set the font for the title
                "fontWeight": "bold",
                "anchor": "start",       # Align the title (start, middle, end)
                # "color": "blue",         # Title color
                "subtitleFontSize": 10,  # Set the subtitle font size
                # "subtitleFont": "Arial", # Set the font for the subtitle
                "subtitleColor": "gray",  # Subtitle color
                # "subtitleFontWeight":'bold'
            },
        },
    }
    return config

def set_theme(theme,
              colors_cat = None,
              colors_cont = None,
              borders = False,
              grid = False
              ):
    if theme == "sci":
        theme_config = theme_sci()

    if borders:
        theme_config['config']["view"]['stroke'] = 'gray'
        theme_config['config']["view"]['strokeWidth'] = .5
        theme_config['config']["axisX"]['domain'] = False
        theme_config['config']["axisX"]['tickOffset'] = 0
        theme_config['config']["axisX"]['ticks'] = False
        theme_config['config']["axisY"]['domain'] = 0
        theme_config['config']["axisY"]['tickOffset'] = 0
        theme_config['config']["axisY"]['ticks'] = False

    if colors_cat is not None:
        theme_config['config']["range"]['category']['scheme'] = colors_cat
    if colors_cont is not None:
        theme_config['config']["range"]['continuous']['scheme'] = colors_cont

    if grid:
        theme_config['config']["axis"]['grid'] = True
        theme_config['config']["axis"]['gridColor'] = 'lightgray'
        theme_config['config']["axis"]['gridOpacity'] = .7
        theme_config['config']["axis"]['gridDash'] = [2,2]
        theme_config['config']["axis"]['gridWidth'] = .5

    @alt.theme.register("new_theme", enable=True)
    def new_theme(): return theme_config
    alt.theme.enable('new_theme')
        
# * others

def tables_build_dict(TABLES, FIGURES, tables_fn="tab-{i+1}", figures_fn="fig-{i+1}"):
    for i, (table, _) in enumerate(TABLES.items()):
        TABLES[table]['label'] = table
        TABLES[table]['fn']    = eval(f"f'{tables_fn}'")
    for i, (figure, _) in enumerate(FIGURES.items()):
        FIGURES[figure]['label'] = figure
        FIGURES[figure]['fn']    = eval(f"f'{figures_fn}'")
    return TABLES, FIGURES

def print_system_info(fn, author=False):
    end_time = time.perf_counter()
    execution_time = end_time - START_TIME
    print('-'*40)
    print(watermark.watermark(custom_time="%a %b %d %Y - %H:%M %p",
                              updated=True))
    print(f"Script: {fn}")
    print(f"Code execution time: {execution_time/60:.4f} minutes")
    # print(f"Operating System: {platform.system()} {platform.platform()}")
    print('')
    if author:
        print(watermark.watermark(current_time=True,
                                  author="Diogo Ferrari", 
                                  website='www.diogoferrari.com'))

    print("System Information:")
    print("-------------------")
    print(watermark.watermark(current_time=True,
                              machine=True, 
                              python=True,
                              # gpu=True,
                              ))
    # Memory information
    if psutil:
        mem = psutil.virtual_memory()
        total_memory_gb = mem.total / (1024 ** 3)
        print("Memory available: {:.2f} GB".format(total_memory_gb))
    else:
        print("Memory info: N/A (psutil not installed)")
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    memory_used_mb = memory_info.rss / (1024 * 1024 * 1024)
    print(f"Memory used: {memory_used_mb:.2f} GB")

    print("\nModules used:")
    print("-------------")

# * done
print('done!')
